# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import datetime
import os
import functools
import logging
import signal

from absl import (flags, app)
from tornado.ioloop import IOLoop, PeriodicCallback
import numpy as np

import coin.flow.subscriber as flow_subscriber
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.base.order_gateway import OrderSide
from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.strategy.hk_hard_arb.products import NativeProducts
from coin.base.param_util import to_list
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.exchange.util.feed_checker import get_default_feed_checker_config
from coin.flow.simple_queue import SimpleQueue
from coin.proto.coin_order_checker_pb2 import AccountCheckerConfig
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.accounting.pnl_balance_printer.factory import get_pnl_balance_printer_cls
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.hk_hard_arb.util import load_clean_json
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import (create_live_env_instance, get_live_env_cls)
from coin.strategy.accounting.pnl_balance_printer.util import reserve_as_proto
from coin.strategy.executor.aggressive.factory import construct_agg_executor
from coin.strategy.executor.aggressive.util import \
    convert_executor_config_from_dict as convert_agg_config
from coin.support.accounting.checker.account_checker import AccountChecker
from coin.support.accounting.checker.config_util import convert_checker_config_from_dict
from coin.exchange.base.order_gateway_config import convert_og_config_from_dict
from coin.exchange.huobi_futures.kr_rest.native_private_client import HuobiFuturesOrderSide
from coin.exchange.okex_v3.kr_rest.native_private_client import OkexFuturesOrderSide
from coin.strategy.tool.fill_log_dumper import FillLogDumper
from coin.strategy.basis_strat.fill_info_writer import FillInfo, FillInfoWriter
from coin.strategy.util.request_util import convert_strategy_request_from_dict
from coin.strategy.accounting.fee import get_fee_rate
from coin.strategy.basis_strat.moving_average import MovingObject
from coin.experimental.yuxuan.model.model import LinearModel
from coin.strategy.basis_strat.model.feature import get_features_from_config
from coin.strategy.basis_strat.model.sampler import get_sampler_from_config
from coin.strategy.basis_strat.base_strategy import BaseStrategy

FLAGS = flags.FLAGS


class DummyPosition(object):
  def __init__(self):
    self.long = self.short = 0


def _gen_arb_products(common, products0, products1, ProductType0, ProductType1):
  if 'exch0_quote_currency' in common and 'exch1_quote_currency' in common:
    exch0_quote_currency = common['exch0_quote_currency']
    exch1_quote_currency = common['exch1_quote_currency']
  else:
    raise ValueError(common)

  products0 = NativeProducts.FromTypedStrProduct(exch0_quote_currency,
                                                 sorted(list(products0.keys())),
                                                 ProductType0)
  products1 = NativeProducts.FromTypedStrProduct(exch1_quote_currency,
                                                 sorted(list(products1.keys())),
                                                 ProductType1)
  assert len(products0.products) == len(products1.products)
  return products0, products1


class BasisLmaggSubStrategy(object):
  def __init__(self, parent, product0, product1, fillinfo_writer, strategy_logger):
    self._parent = parent
    self._product0 = product0
    self._product1 = product1
    self._fillinfo_writer = fillinfo_writer
    self._strategy_logger = strategy_logger
    self._ticker0 = product0.full_symbol
    self._ticker1 = product1.full_symbol
    self._tick_info1 = (get_holder_from_product(product1).product_info.gen_tick(),
                        get_holder_from_product(product1).product_info.gen_tick_qty())
    if hasattr(self._tick_info1[0], 'tick_size'):
      ticksize = self._tick_info1[0].tick_size
      self._next_price_func = lambda p: p + ticksize - (p % ticksize)
      self._prev_price_func = lambda p: p - ((p % ticksize) or ticksize)
    else:
      self._next_price_func = self._tick_info1[0].get_next_price
      self._prev_price_func = self._tick_info1[0].get_prev_price

    self.og1 = None
    self._agg_executor1 = None
    self.logger = parent.logger
    self._last_feed = parent._last_feed

    self._pos1 = 0
    self._pos1_ts = 0
    self._og_pos1 = 0
    self._og_pos1_ts = 0
    self._update_pos_with_balance = FLAGS.update_pos_with_balance
    self._update_position_duration = FLAGS.update_pos_with_balance_sec * 10**9
    self._update_working_qty_duration = FLAGS.update_working_qty_sec * 10**9
    self._equal_theshold = None

    self._lot_size = None
    self._max_pos = None
    self._min_pos = None
    self._agg_taker_fee = None
    self._expected_working_buy = 0
    self._expected_working_sell = 0
    self._last_submit_fill_ts = 0
    # Handle agg working order update delay.
    self._agg_working_order_fills = dict()

    # Decide if book has changed.
    self._book_change_levels = min(FLAGS.book_change_levels, FLAGS.book_levels)

    self._get_pos_func = None
    self._is_buy_order_func = None
    self._dummy_position = DummyPosition()
    self._raw_run_size = 1.0
    self._relative_run_size = None

    # Aggressive orders related
    self._force_fire_cooldown = FLAGS.force_fire_cooldown_sec * 1e9
    self._force_fire_cooldown_window = FLAGS.force_fire_cooldown_window_sec * 1e9
    self._force_fire_max_count = FLAGS.force_fire_max_count
    self._force_fire_qty_mul = FLAGS.force_fire_qty_multiplier
    self._force_fire_cap_qty = self._lot_size
    self._force_fire_bot_qty = 0
    self._aggressive_amt = FLAGS.aggressive_bp * 1e-4
    self._aggressive_enabled = abs(self._aggressive_amt) > 1e-6
    self._agg_buys = MovingObject(self._force_fire_cooldown_window)
    self._agg_sells = MovingObject(self._force_fire_cooldown_window)
    # Model and features
    self._models = []
    self._features = None
    self._sampler = None
    # Params
    self._threshold_mul = FLAGS.threshold_multiplier
    self._min_predict_threshold = FLAGS.min_predict_bp * 1e-4

  def ready(self):
    # 30s warm up time.
    return self.get_ts() - self._parent._start_time > 30e9

  def get_mid_price(self, product):
    bids, asks = self._last_feed[product]
    return (bids[0][0] + asks[0][0]) / 2.0

  def register_model(self, model):
    self._models.append(model)

  def register_sampler(self, sampler):
    self._sampler = sampler
    if sampler:
      self.logger.info('Sampler %s' % sampler)

  def register_features(self, features):
    self._features = features

  def _init_get_pos_func(self, og):
    product_type = self._product1.product_type
    exchange = self._product1.exchange
    if product_type == 'Futures':
      if exchange == 'Huobi':
        self._get_pos_func = lambda: \
            (og.get_huobi_futures_position(self._product1) or self._dummy_position).long - \
            (og.get_huobi_futures_position(self._product1) or self._dummy_position).short
        self._is_buy_order_func = lambda order_side: order_side in [
            HuobiFuturesOrderSide.BUY_OPEN, HuobiFuturesOrderSide.BUY_CLOSE, OrderSide.BUY
        ]
      elif exchange == 'Kraken':
        self._get_pos_func = lambda: \
            getattr(og.og_info.get_product_position_proto(self._product1), 'net_position', 0.0)
        self._is_buy_order_func = lambda order_side: order_side == OrderSide.BUY
      elif exchange == 'Okex':
        self._get_pos_func = lambda: \
            (og.get_okex_futures_position(self._product1) or self._dummy_position).long - \
            (og.get_okex_futures_position(self._product1) or self._dummy_position).short
        self._is_buy_order_func = lambda order_side: order_side in \
            [OkexFuturesOrderSide.BUY_OPEN, OkexFuturesOrderSide.BUY_CLOSE]
      elif exchange in ['Bitmex', 'Binance', 'Bitflyer']:
        self._get_pos_func = lambda: og.get_position(self._product1)
        self._is_buy_order_func = lambda order_side: order_side == OrderSide.BUY
    elif product_type == 'CurrencyPair':
      self._get_pos_func = lambda: og.get_balance(self._product1.base.currency)
      self._is_buy_order_func = lambda order_side: order_side == OrderSide.BUY
    assert self._get_pos_func is not None, '%s %s' % (product_type, exchange)

  def write_executor_config(self):
    if self._config1_proto is not None:
      self._strategy_logger.write_executor_config(self._config1_proto)

  def _init_model_features(self, config):
    assert 'model_name' in config
    assert 'feature_name' in config
    model_path = os.path.join(FLAGS.model_dir,
                              FLAGS.strategy_name,
                              '%s.pickle' % config['model_name'])
    feature_path = os.path.join(FLAGS.model_dir,
                                FLAGS.strategy_name,
                                '%s.json' % config['feature_name'])
    assert os.path.exists(model_path), model_path
    assert os.path.exists(feature_path), feature_path
    config = load_clean_json(feature_path)
    features = config['features']
    sampler = config.get('sampler', None)
    model = LinearModel()
    model.load_model(model_path)
    features = get_features_from_config(features, self, [self._product0, self._product1])
    sampler = get_sampler_from_config(sampler, self, [self._product0, self._product1])
    self.register_model(model)
    self.register_features(features)
    self.register_sampler(sampler)
    self.logger.info('Model and features inited.')

  def init_executors(self, req1, config1, og1):
    self._init_model_features(config1)
    self._relative_run_size = 1.0
    self.og1 = og1
    self._raw_config1 = config1.copy()
    self._config1 = config1.copy()
    self._config1['force_fire'] = True

    if FLAGS.run_size != 1.0:
      assert 0 < FLAGS.run_size < 1, FLAGS.run_size
      self._raw_run_size = FLAGS.run_size
      max_pos = self._config1['max_pos']
      min_pos = self._config1['min_pos']
      neutral_pos = (max_pos + min_pos) / 2.0
      max_exposure = (max_pos - min_pos) / 2.0 * FLAGS.run_size
      self._config1['max_pos'] = neutral_pos + max_exposure
      self._config1['min_pos'] = neutral_pos - max_exposure
      self._config1['lot_size'] = self._config1['lot_size'] * FLAGS.run_size
      self.logger.info('Trimmed size %f %f' % (self._config1['min_pos'], self._config1['max_pos']))

    self._config1_proto = convert_agg_config(self._product1.symbol, self._config1)
    self._agg_executor1 = construct_agg_executor(req1, self._config1_proto, self.og1, self.logger)

    self._init_get_pos_func(og1)
    self._equal_theshold = max(self._config1['lot_size'] * 1e-4, 1e-5)
    self._lot_size = self._config1['lot_size']
    self._force_fire_cap_qty = self._lot_size
    assert FLAGS.order_qty_ratio_threshold < 0.5
    self._force_fire_bot_qty = self._lot_size * FLAGS.order_qty_ratio_threshold
    self._max_pos = self._config1['max_pos']
    self._min_pos = self._config1['min_pos']
    self._threshold_mul = self._config1.get('threshold_multiplier', self._threshold_mul)

    _, self._agg_taker_fee = get_fee_rate(self.og1.market_type, self.og1.exchange, self._product1)
    self._raw_agg_taker_fee = self._agg_taker_fee
    if hasattr(self._product1, 'expiry') and self._product1.expiry:
      self.logger.info('Futures product, set fees to 0, %s' % self._product1.expiry)
      self._agg_taker_fee = 0.0
    self._update_pos1(self._get_pos_func())
    self._og_pos1 = self._pos1
    self._og_pos1_ts = self._get_ts()
    assert abs(self._agg_taker_fee) <= 1e-2, self._agg_taker_fee
    self.logger.info('%s init pos1 %f' % (self._ticker0, self._pos1))

  def update_runsize_if_necessary(self, run_size):
    if self._relative_run_size is None or \
            np.isclose(self._relative_run_size, run_size):
      return
    self._relative_run_size = run_size
    overall_run_size = self._raw_run_size * self._relative_run_size
    self._config1 = self._raw_config1.copy()
    self._config1['force_fire'] = True
    max_pos = self._config1['max_pos']
    min_pos = self._config1['min_pos']
    neutral_pos = (max_pos + min_pos) / 2.0
    max_exposure = (max_pos - min_pos) / 2.0 * overall_run_size
    self._config1['max_pos'] = neutral_pos + max_exposure
    self._config1['min_pos'] = neutral_pos - max_exposure
    self._config1['lot_size'] = self._config1['lot_size'] * overall_run_size
    self._equal_theshold = max(self._config1['lot_size'] * 1e-4, 1e-5)
    self._lot_size = self._config1['lot_size']
    self._force_fire_cap_qty = self._lot_size * self._force_fire_qty_mul
    assert FLAGS.order_qty_ratio_threshold < 0.5
    self._force_fire_bot_qty = self._lot_size * FLAGS.order_qty_ratio_threshold
    self._max_pos = self._config1['max_pos']
    self._min_pos = self._config1['min_pos']
    self.logger.info('Update run_size %f %f %f %f' % (
        self._relative_run_size, self._lot_size, self._config1['min_pos'],
        self._config1['max_pos']))
    self._config1_proto = convert_agg_config(self._product1.symbol, self._config1)
    self._agg_executor1.set_config(self._config1_proto)

  def _update_pos1(self, pos):
    self._pos1 = pos
    self._pos1_ts = self._get_ts()

  def _get_ts(self):
    return self._parent._ts

  def get_ts(self):
    return self._get_ts()

  def _is_book_valid(self):
    return self._get_ts() - self._book_ts[0] < self._book_invalid_period

  def _position_equal(self, p1, p2):
    return abs(p1 - p2) < self._equal_theshold

  def _get_working_qty(self):
    # Use working orders to update |self._expected_working_qty|.
    if self._get_ts() - self._last_submit_fill_ts > self._update_working_qty_duration:
      order_id_set = set()
      working_buy = 0.0
      working_sell = 0.0
      for order in self._agg_executor1._wo.buy_working_orders:
        working_buy += (order.qty - self._agg_working_order_fills.get(order.order_id, 0.0))
        order_id_set.add(order.order_id)
      for order in self._agg_executor1._wo.sell_working_orders:
        working_sell += (order.qty - self._agg_working_order_fills.get(order.order_id, 0.0))
        order_id_set.add(order.order_id)
      # Remove order ids that are not in working orders.
      obsolete_ids = self._agg_working_order_fills.keys() - order_id_set
      for order_id in obsolete_ids:
        del self._agg_working_order_fills[order_id]
      if not self._position_equal(self._expected_working_buy, working_buy):
        self.logger.info('working qty expected %f, updated %f' %
                         (self._expected_working_buy, working_buy))
      if not self._position_equal(self._expected_working_sell, working_sell):
        self.logger.info('working qty expected %f, updated %f' %
                         (self._expected_working_sell, working_sell))
      self._expected_working_buy = working_buy
      self._expected_working_sell = working_sell
    return self._expected_working_buy, self._expected_working_sell

  def _is_pos_valid(self):
    ts = self._get_ts()
    return ts - max(self._pos1_ts, self._og_pos1_ts) > self._update_position_duration

  def _get_balance(self):
    og_pos = self._get_pos_func()
    if self._update_pos_with_balance:
      if not self._position_equal(og_pos, self._og_pos1):
        self._og_pos1 = og_pos
        self._og_pos1_ts = self._get_ts()
      if not self._position_equal(self._pos1, self._og_pos1) \
              and self._is_pos_valid():
        self.logger.info('Update pos1 with og_pos1, %f %f' % (self._pos1, self._og_pos1))
        self._pos1 = self._og_pos1
    else:
      if self._position_equal(self._pos1, og_pos):
        self._pos1 = og_pos
    return self._pos1

  def _get_agg_working_qty(self, mo):
    qty = 0.0
    for ts, obj in mo.iteritems():
      qty += obj[2]
    return qty

  def _get_force_buy_qty(self, agg_price, asks):
    self._agg_buys.update(self._get_ts())
    size = self._agg_buys.size()
    qty = 0
    working_qty = 0
    avail_qty = 0
    if size >= self._force_fire_max_count:
      return 0
    elif size > 0:
      last_ts = self._agg_buys.get_latest_ts()
      if self._get_ts() - last_ts < self._force_fire_cooldown:
        return 0
      last_agg_price, last_ba, _ = \
          self._agg_buys.get_latest_data()
      # Order book not updated and agg price not improved.
      if np.isclose(last_ba[0], asks[0][0]) \
              and np.isclose(last_ba[1], asks[0][1]):
        return 0
      working_qty = self._get_agg_working_qty(self._agg_buys)
    for p, q in asks:
      if p > agg_price + 1e-5:
        break
      avail_qty += q
    vacancy_qty = self._max_pos - self._get_balance() - working_qty
    qty = min(vacancy_qty, self._force_fire_cap_qty, avail_qty)
    return qty

  def _get_force_sell_qty(self, agg_price, bids):
    self._agg_sells.update(self._get_ts())
    size = self._agg_sells.size()
    qty = 0
    working_qty = 0
    avail_qty = 0
    if size >= self._force_fire_max_count:
      return 0
    elif size > 0:
      last_ts = self._agg_sells.get_latest_ts()
      if self._get_ts() - last_ts < self._force_fire_cooldown:
        return 0
      last_agg_price, last_bb, _ = \
          self._agg_sells.get_latest_data()
      # Order book not updated and agg price not improved.
      if np.isclose(last_bb[0], bids[0][0]) \
              and np.isclose(last_bb[1], bids[0][1]):
        return 0
      working_qty = self._get_agg_working_qty(self._agg_sells)
    for p, q in bids:
      if p < agg_price - 1e-5:
        break
      avail_qty += q
    vacancy_qty = self._get_balance() - self._min_pos - working_qty
    qty = min(vacancy_qty, self._force_fire_cap_qty, avail_qty)
    return qty

  def _process(self):
    featureslist = []
    for feature in self._features:
      featureslist += feature.get_features()
    bids1, asks1 = self._last_feed[self._product1]
    features = np.array([featureslist])
    preds = [model.predict(features)[0] for model in self._models]
    mean_pred = np.mean(preds)
    spread = (asks1[0][0] - bids1[0][0]) / bids1[0][0]
    threshold = (spread + 2 * self._raw_agg_taker_fee) * self._threshold_mul
    order_qty = 0
    cancel_all_existing_buy = False
    cancel_all_existing_sell = False
    if mean_pred > max(threshold, self._min_predict_threshold):
      order_side = OrderSide.BUY
      cancel_all_existing_sell = True
      if self._aggressive_enabled:
        order_price = self._prev_price_func(asks1[0][0] * (1 + self._aggressive_amt))
      else:
        order_price = asks1[0][0]
      order_qty = self._get_force_buy_qty(order_price, asks1)
    elif mean_pred < min(-threshold, -self._min_predict_threshold):
      order_side = OrderSide.SELL
      cancel_all_existing_buy = True
      if self._aggressive_enabled:
        order_price = self._next_price_func(bids1[0][0] * (1 - self._aggressive_amt))
      else:
        order_price = bids1[0][0]
      order_qty = self._get_force_sell_qty(order_price, bids1)
    if FLAGS.dry_run:
      return
    if order_qty > self._force_fire_bot_qty + 1e-6:
      ret = self._agg_executor1.submit_agg_order(order_side=order_side,
                                                 order_price=order_price,
                                                 order_qty=abs(order_qty))
      if ret:
        if order_side == OrderSide.BUY:
          self._expected_working_buy += order_qty
          self._agg_buys.update(self.get_ts(), (order_price, asks1[0], order_qty))
        else:
          self._expected_working_sell += order_qty
          self._agg_sells.update(self.get_ts(), (order_price, bids1[0], order_qty))
        self._last_submit_fill_ts = self._get_ts()
        self.logger.info('agg submitted return code %s' % ret)
        self._fillinfo_writer.write(
            FillInfo('agg submitted',
                     symbol=self._product1.symbol,
                     exchange=self._product1.exchange,
                     fill_price=order_price,
                     fill_qty=order_qty,
                     order_side=order_side,
                     timestamp=self._get_ts()))
        self.logger.info('features: %s, %s' % (featureslist, mean_pred))
    if cancel_all_existing_buy:
      self._agg_executor1.cancel_all_existing_buy()
    if cancel_all_existing_sell:
      self._agg_executor1.cancel_all_existing_sell()

  def on_book(self, product, book):
    if not self._agg_executor1:
      return

    for feature in self._features:
      feature.on_book(product, book)

    if not self.ready():
      return

    if self._sampler:
      self._sampler.on_book(product, book)
      if not self._sampler.sample():
        self._agg_executor1.manage_agg_orders()
        return

    self._process()
    self._agg_executor1.manage_agg_orders()

  def on_trade(self, product, trade):
    if self._features:
      for feature in self._features:
        feature.on_trade(product, trade)

  def on_fill(self, fill_price, fill_qty, order):
    self._agg_working_order_fills[order.order_id] = \
        self._agg_working_order_fills.get(order.order_id, 0.0) + fill_qty
    self._fillinfo_writer.write(
        FillInfo('Filled',
                 symbol=order.product.symbol,
                 exchange=order.product.exchange,
                 fill_price=fill_price,
                 fill_qty=fill_qty,
                 order_side=order.order_side,
                 timestamp=self._get_ts()))
    if self._is_buy_order_func(order.order_side):
      self._update_pos1(self._pos1 + fill_qty * (1 - self._agg_taker_fee))
    else:
      self._update_pos1(self._pos1 - fill_qty)
    self._last_submit_fill_ts = self._get_ts()
    self.logger.info('expected %s pos, %f' % (self._ticker0, self._pos1))


class BasisLmaggStrategy(BaseStrategy):
  def __init__(self, strategy_params):
    self.logger = logging.getLogger(__name__)
    self._config = strategy_params['config']
    self._exch0_subreq = SubscriptionRequest.from_dict(
        BaseStrategy.update_sub_config(self._config['common']['exch0_subreq']))
    self._exch1_subreq = SubscriptionRequest.from_dict(
        BaseStrategy.update_sub_config(self._config['common']['exch1_subreq']))

    self._exchange0 = self._exch0_subreq.exchange
    self._exchange1 = self._exch1_subreq.exchange
    self._og_config = strategy_params['og_config']
    self._balance_checker = None
    account_checker_config = self._get_account_checker_config()
    if account_checker_config is not None:
      self._balance_checker = AccountChecker(account_checker_config)

    self._book_builder_clses = (get_flow_book_builder_cls(self._exch0_subreq),
                                get_flow_book_builder_cls(self._exch1_subreq))

    self._live_env_clses1 = get_live_env_cls(self._exch1_subreq)
    self._live_env1 = None
    self.og1 = None

    self._products0, self._products1 = _gen_arb_products(self._config['common'],
                                                         self._config['products0'],
                                                         self._config['products1'],
                                                         self._book_builder_clses[0].ProductType,
                                                         self._book_builder_clses[1].ProductType)
    self._is_future = self._products1.products[0].product_type == 'Futures'

    exch0_exchange_feed_checker = self._config['common'].get('exch0_exchange_feed_checker', False)
    exch1_exchange_feed_checker = self._config['common'].get('exch1_exchange_feed_checker', False)
    feed_checker_products = []
    if exch0_exchange_feed_checker:
      feed_checker_products += self._products0.products
    if exch1_exchange_feed_checker:
      feed_checker_products += self._products1.products
    feed_checker_config = get_default_feed_checker_config(feed_checker_products)

    multi_feed_subsystem_cls = BaseStrategy.get_multi_feed_subsystem()
    self._feed = multi_feed_subsystem_cls(feed_checker_config=feed_checker_config)
    self._real_time = datetime.datetime.now()
    self._start_time = int(float(self._real_time.strftime('%s.%f')) * 1e9)

    expiry = BasisLmaggStrategy._get_config_key(self._products0.products[0])
    self._ref_weekly = 'THIS_WEEK' in expiry
    self._ref_quarter = 'QUARTER' in expiry
    self._run_size = 1.0
    self.logger.info('%s, ref weekly: %s, ref quarter: %s' %
                     (expiry, self._ref_weekly, self._ref_quarter))

    self._ts = None
    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None
    self._balance_printer = None
    self._fillinfo_writer = FillInfoWriter('fills.csv', FLAGS.strategy_name)
    self._fill_dumper = FillLogDumper(self._exch1_subreq)

    self._book_levels = FLAGS.book_levels
    self._last_feed = {}
    self._book_map = BookMap(self._products0.products + self._products1.products,
                             logger=self.logger)

    self._reserve_map = self._get_reserve_map()
    self._strategy_request = convert_strategy_request_from_dict(
        self._config['common']['strategy_request'])
    self._reserve_proto = reserve_as_proto(self._exch1_subreq, self._reserve_map)
    self._strategy_logger = StrategyLogProvider(self._strategy_request)

    self._substrategies = {}
    self._usubstrategies = []
    for product0, product1 in zip(self._products0.products, self._products1.products):
      substrategy = BasisLmaggSubStrategy(self,
                                          product0,
                                          product1,
                                          self._fillinfo_writer,
                                          self._strategy_logger)
      self._substrategies[product0] = self._substrategies[product1] = substrategy
      self._usubstrategies.append(substrategy)

  def periodic_write_logs(self):
    self.og1.og_logger.write_account_request()
    self._strategy_logger.write_request()
    self._strategy_logger.write_reserve(self._reserve_proto)
    for strat in self._usubstrategies:
      strat.write_executor_config()
    self._ioloop.add_timeout(datetime.timedelta(hours=1), self.periodic_write_logs)

  @property
  def strategy_logger(self):
    return self._strategy_logger

  def _update_prices(self, product, book):
    bids = book.get_bid_array(self._book_levels)
    asks = book.get_ask_array(self._book_levels)
    feed = self._last_feed.get(product, None)
    new_feed = (bids, asks)
    if feed is None or not (np.all(np.isclose(feed[0], new_feed[0]))
                            and np.all(np.isclose(feed[1], new_feed[1]))):
      self._last_feed[product] = new_feed
      return True
    return False

  def on_book(self, product, book):
    self._ts = book.timestamp
    changed = self._update_prices(product, book)
    if changed:
      self._substrategies[product].on_book(product, book)
    self._book_map.set_book(product, book)

  def on_trade(self, product, trade):
    self._ts = trade.timestamp
    self._substrategies[product].on_trade(product, trade)

  def on_fill(self, fill_price, fill_qty, order):
    self._substrategies[order.product].on_fill(fill_price, fill_qty, order)
    self._fill_dumper.dump_fill(fill_price, fill_qty, order)

  def _get_account_checker_config(self):
    if self._config.get('account_checker') is None:
      return None

    config_keys = {
        'reserve',
        'lot_size',
        'stack_size',
        'ratio_min_balance',
        'ratio_max_balance',
        'ratio_min_hold_balance',
    }

    balance_check_config_list = []
    for currency, config in self._config['account_checker'].items():
      config = {**config}
      if 'lot_size' not in config:
        for product in self._config['products']:
          base = SharedProduct.FromStrNativeProduct(product).base.currency
          if base == currency:
            config['lot_size'] = self._config['products'][product]['lot_size']
            break

      assert config_keys == set(config.keys()), config
      reserve = config['reserve']
      config['currency'] = currency
      config['min_balance'] = reserve * config['ratio_min_balance']
      config['max_balance'] = reserve * config['ratio_max_balance']
      config['min_available_balance'] = 0.0
      config['max_available_balance'] = config['lot_size'] * config['stack_size']
      config['min_hold_balance'] = reserve * config['ratio_min_hold_balance']
      config['max_hold_balance'] = reserve * config['ratio_max_balance']
      currency_checker_config = convert_checker_config_from_dict(config)
      balance_check_config_list.append(currency_checker_config)

    return AccountCheckerConfig(each_balance_check_config=balance_check_config_list)

  def signal_handler(self, sig, frame):
    self.logger.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self.on_exit)

  def on_exit(self):
    self._exit_attempt += 1
    if self._exit_attempt > 5:
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None and not FLAGS.dry_run:
      ogs = to_list([self.og1])
      self.logger.info('Exit callback executed.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._feed.start(ioloop, queue, flow_sub)
    self._live_env1.start()
    self.start_print_states()

  def subscribe_feeds(self):
    func = self.on_book
    for product in self._products0.products:
      feed_req = FeedSubscriptionRequest.create([product], self._exch0_subreq)
      self._feed.subscribe_book(feed_req, functools.partial(func, product))
      self._feed.subscribe_trade(feed_req, functools.partial(self.on_trade, product))
    for product in self._products1.products:
      feed_req = FeedSubscriptionRequest.create([product], self._exch1_subreq)
      self._feed.subscribe_book(feed_req, functools.partial(func, product))
      self._feed.subscribe_trade(feed_req, functools.partial(self.on_trade, product))

  def subscribe_live_envs(self, ioloop, queue, flow_sub):
    og_params = self._config.get('og_params', {})
    og_params1 = convert_og_config_from_dict(self._exch1_subreq, og_params)

    self._live_env1 = create_live_env_instance(self._live_env_clses1,
                                               self._products1.currencies,
                                               self._products1.products,
                                               self._og_config[0],
                                               og_params=og_params1,
                                               ioloop=ioloop,
                                               queue=queue,
                                               flow_sub=flow_sub,
                                               og_log_callback=self._strategy_logger.write_og_log)
    self._live_env1.on_order_gateway_reset_callback.append(self.register_order_gateway)

  def _get_reserve_map(self):
    reserves = dict()
    for product in self._products1.products:
      key = BasisLmaggStrategy._get_config_key(product)
      reserves[product.base.currency] = (self._config['products1'][key]['max_pos']
                                         + self._config['products1'][key]['min_pos']) / 2.0
    return reserves

  def register_balance_printer(self, pos_bias_map):
    printer_cls = get_pnl_balance_printer_cls(self._exch1_subreq)
    if len(pos_bias_map) > 0:
      self._balance_printer = printer_cls(self._products1.products,
                                          self._book_map,
                                          self._reserve_map,
                                          logger=self.logger,
                                          pos_bias_map=pos_bias_map)
    else:
      self._balance_printer = printer_cls(self._products1.products,
                                          self._book_map,
                                          self._reserve_map,
                                          logger=self.logger)
    self._balance_printer.register_order_gateway(self.og1)
    PeriodicCallback(self.dump_pnl_balance, 30 * 1000).start()

  def start_print_states(self):
    for strat in self._usubstrategies:
      # TODO print states of strategy.
      break
    self._ioloop.add_timeout(datetime.timedelta(seconds=10), self.start_print_states)

  def dump_pnl_balance(self, periodic=False):
    if self._balance_printer is not None:
      self._balance_printer.print_pnl_balance()
      pnl_balances = self._balance_printer.calculate_pnl_balance()
      if pnl_balances is not None:
        for pnl_balance in pnl_balances:
          self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance,
                                                  quote=None,
                                                  subreq_id=None)
    self._update_runsize_if_necessary()

  def _update_runsize_if_necessary(self):
    self._real_time = datetime.datetime.now()
    run_size = BaseStrategy.check_rollover(self._ref_weekly, self._ref_quarter, self._real_time)
    schedules = BaseStrategy.read_maintenance_as_json(FLAGS.maintenance_file,
                                                      [self._exchange0, self._exchange1])
    for p0, p1 in zip(self._products0.products, self._products1.products):
      run_size2 = min(BaseStrategy.check_maintenance(schedules, p0, self._real_time),
                      BaseStrategy.check_maintenance(schedules, p1, self._real_time))
      run_size2 = min(run_size, run_size2)
      self._substrategies[p1].update_runsize_if_necessary(run_size2)
    if schedules:
      self.logger.info('Schedules count %s' % len(schedules))

  @staticmethod
  def _get_config_key(product):
    key = product.subscription_symbol
    if product.exchange == 'Huobi' and product.product_type == 'Futures':
      key = key.replace('CURRENT', 'THIS')
      key = key.replace('THIS_QUARTER', 'QUARTER')
    return key

  def _get_pos_bias_map(self):
    pos_bias_map = dict()
    for prod1 in self._products1.products:
      key = BasisLmaggStrategy._get_config_key(prod1)
      config = self._config['products1'][key]
      if 'pos_bias' in config:
        assert 'pos_bias_price' in config
        pos_bias_map[prod1] = (config['pos_bias'], config['pos_bias_price'])
    self.logger.info('pos_bias_map %s' % pos_bias_map)
    return pos_bias_map

  def register_order_gateway(self, name, order_gateway):
    self.logger.info('OG initialized: %s', name)
    if order_gateway.exchange == self._exchange1:
      self.og1 = order_gateway
    # cancel working orders first after OGs are ready.
    if not self._order_canceled_after_og_ready:
      order_gateway.cancel_all()
      self.logger.info('Orders canceled after OGs are ready.')
      self._order_canceled_after_og_ready = True
    self._exit_executor = SimpleCancelAllExecutor()

    if self.og1:
      pos_bias_map = self._get_pos_bias_map()
      self.register_balance_printer(pos_bias_map)
      self.og1.register_fill_callback(self.on_fill)
      for prod0, prod1 in zip(self._products0.products, self._products1.products):
        key = BasisLmaggStrategy._get_config_key(prod1)
        config1 = self._config['products1'][key].copy()
        config1.update(self._config.get('agg_executor_configs', {}))
        self._substrategies[prod1].init_executors(self._exch1_subreq, config1, self.og1)
      self.periodic_write_logs()


def get_param_set(og_configs, config_filename):
  og_config = og_configs.split(',')
  config = load_clean_json(config_filename)
  return {
      'og_config': og_config,
      'config': config,
  }


def main(argv):
  assert FLAGS.strategy_name
  assert FLAGS.model_dir

  strategy_params = get_param_set(og_configs=FLAGS.og_configs,
                                  config_filename=FLAGS.config_filename)

  queue = SimpleQueue('basisunhedged')
  flow_sub = flow_subscriber.from_queue(queue)
  ioloop = IOLoop.current()

  strategy = BasisLmaggStrategy(strategy_params)
  strategy.subscribe_feeds()
  strategy.subscribe_live_envs(ioloop, queue, flow_sub)
  strategy.start(ioloop, queue, flow_sub)

  # exit_after = datetime.timedelta(days=1)
  # ioloop.add_timeout(exit_after, strategy.on_exit)

  signal.signal(signal.SIGTERM, strategy.signal_handler)
  signal.signal(signal.SIGINT, strategy.signal_handler)

  ioloop.start()


if __name__ == '__main__':
  flags.DEFINE_string('og_configs', '', '')

  flags.DEFINE_string('config_filename', '', '')

  flags.DEFINE_boolean('dry_run', False, '')

  flags.DEFINE_float('book_invalid_sec', 3.5, '')

  flags.DEFINE_boolean('update_pos_with_balance', False, '')

  flags.DEFINE_float('update_pos_with_balance_sec', 4, '')

  flags.DEFINE_float('update_working_qty_sec', 3, '')

  flags.DEFINE_integer('book_levels', 5, '')

  flags.DEFINE_integer('book_change_levels', 3, '')

  flags.DEFINE_string('strategy_name', '', '')

  flags.DEFINE_float('run_size', 1.0, '')

  flags.DEFINE_string('model_dir', '', '')

  flags.DEFINE_float('threshold_multiplier', 1.5, '')

  flags.DEFINE_float('force_fire_cooldown_window_sec', 5, '')

  flags.DEFINE_float('force_fire_cooldown_sec', 1, '')

  flags.DEFINE_integer('force_fire_max_count', 3, '')

  flags.DEFINE_float('force_fire_qty_multiplier', 1.5, '')

  flags.DEFINE_float('aggressive_bp', 0.5, '')

  flags.DEFINE_float('min_predict_bp', 0.5, '')

  flags.DEFINE_float('order_qty_ratio_threshold', 0, 'Only fire when qty > lot_size * threshold.')

  flags.DEFINE_string('maintenance_file', 'coin_deploy/basis_strat/schedule/maintenance.json', '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
